const fs = require('fs');
const path = require('path');
const xml2js = require('xml2js');
const labelDef = require('./Def');
const BaseLabelLoader = require('../../BaseLabelLoader');
const { getImageSize } = require('../../../utils/imageUtils');

class VOCLabelLoader extends BaseLabelLoader {
    // 加载数据
    async load(label_info) {
        // console.log('VOCLabelLoader::load label_info = ', label_info);
        try {
            const xmlData = fs.readFileSync(label_info.label_path, 'utf8');
            const xmlRoot = await xml2js.parseStringPromise(xmlData);
            const annotation = xmlRoot.annotation;
            // const segmented = annotation.segmented[0];

            let result = [];
            if (annotation.object) {
                for (let i = 0; i < annotation.object.length; i++) {
                    let jsonObj = labelDef();
                    jsonObj.label_format = 'voc';
                    jsonObj.id = i;

                    let object = annotation.object[i];
                    jsonObj.category_id = null; // TODO： 看下是否需要反查下，暂时先不增加
                    jsonObj.category = object.name ? object.name[0] : 'unknown';
                    jsonObj.pose = object.pose ? object.pose[0] : 'unspecified';
                    jsonObj.truncated = object.truncated ? parseInt(object.truncated[0], 10) : 0;
                    jsonObj.occluded = object.occluded ? parseInt(object.occluded[0], 10) : 0;
                    jsonObj.difficult = object.difficult ? parseInt(object.difficult[0], 10) : 0;

                    const bndbox = object.bndbox[0];
                    let xmin = parseInt(bndbox.xmin[0]);
                    let ymin = parseInt(bndbox.ymin[0]);
                    let xmax = parseInt(bndbox.xmax[0]);
                    let ymax = parseInt(bndbox.ymax[0]);
                    let area = (xmax - xmin) * (ymax - ymin);

                    jsonObj.area = area > 0 ? area : 0;
                    jsonObj.bbox = [xmin, ymin, xmax, ymax];
                    result.push(jsonObj);
                }
            }
            return result;
        } catch (error) {
            console.error('Error load label ', label_info, '\nerror: ', error);
            return { code: -1, msg: error.message };
        }
    }

    // 保存数据
    async save(label_info) {
        // console.log('VOCLabelLoader::save label_info = ', label_info);
        // 转换为voc xml格式
        let vocXmlPath = label_info.label_url;
        let vocXml = await this.fromStudio(label_info);
        // console.log(vocXml);
        const builder = new xml2js.Builder();
        const vocXmlString = builder.buildObject(vocXml);
        fs.writeFileSync(vocXmlPath, vocXmlString);

        // 更新Datasets和ListDB的信息
        let dataset_id = label_info.dataset_id;
        this.app.db && this.app.db.datasets.updateDataset({ id: dataset_id });
        let list_db = this.app.list_dbs[dataset_id];
        list_db && list_db.pairs.updateLabelPath(label_info.pairs_id);
    }

    async toString(label_info, categories = null) {
        let vocXml = await this.fromStudio(label_info, categories);
        const builder = new xml2js.Builder();
        const vocXmlString = builder.buildObject(vocXml);
        return vocXmlString;
    }

    // 转换数据
    async fromStudio(label_info, categories = null) {
        // console.log('VOCLabelLoader::fromStudio label_info = ', label_info);
        let dataset_id = label_info.dataset_id;
        let labels = label_info.labels;
        let image_url = label_info.image_url;

        // 读取图片的宽高，更新到voc的xml里面
        const { width: imageWidth, height: imageHeight } = getImageSize(image_url);

        // 保存voc xml格式
        let xmlRoot = {
            annotation: {
                filename: path.basename(image_url),
                size: {
                    width: imageWidth,
                    height: imageHeight,
                    depth: 3
                },
                segmented: 0,
                object: []
            }
        };

        // 获取categories信息
        if (!categories) {
            categories = await this.app.db.categories.getCategories(dataset_id);
        }
        for (let i = 0; i < labels.length; i++) {
            if (!labels[i]) continue;
            let jsonObj = labels[i].data;
            let xmlObject = {};

            // 处理category
            // const category = categories.find((el) => el.category_name === jsonObj.category);
            // if (category) {
            //     xmlObject.name = jsonObj.category;
            // }
            xmlObject.name = jsonObj.category;

            // 处理bbox
            xmlObject.bndbox = {
                xmin: Math.round(jsonObj.bbox[0]),
                ymin: Math.round(jsonObj.bbox[1]),
                xmax: Math.round(jsonObj.bbox[2]),
                ymax: Math.round(jsonObj.bbox[3])
            };

            // 处理通用的属性
            if (jsonObj) {
                let fields = ['truncated', 'difficult', 'pose'];
                for (let field of fields) {
                    xmlObject[field] = jsonObj[field];
                }
            }

            xmlRoot.annotation.object.push(xmlObject);
        }
        return xmlRoot;
    }
}

module.exports = VOCLabelLoader;
